from linecache import getline
import numpy as np


# 洪水填充算法
# c: 地图中洪水起始点x值；
# r:地图中洪水起始点y值；
# mask：地图中所有高程小于70m的遮盖
def floodFill(c, r, mask):
    # 已填充的单元格
    filled = set()
    # 待填充的单元格
    fill = set()
    fill.add((c, r))
    width = mask.shape[1] - 1
    height = mask.shape[0] - 1
    # 输出的淹没数组
    flood = np.zeros_like(mask, dtype=np.int8)
    print(mask)
    while fill:
        # 获取单元格
        x, y = fill.pop()
        if y == height or x == width or x < 0 or y < 0:
            # 不填充
            continue
        if mask[y][x] == 1:
            # 填充
            flood[y][x] = 1
            filled.add((x, y))
            # 检查相邻单元格
            west = (x - 1, y)
            east = (x + 1, y)
            north = (x, y - 1)
            south = (x, y + 1)
            if west not in filled:
                fill.add(west)
            if east not in filled:
                fill.add(east)
            if north not in filled:
                fill.add(north)
            if south not in filled:
                fill.add(south)
    return flood


# 创建洪水
source = "Grid/FloodFill/terrain.asc"
output = "Grid/FloodFill/flood.asc"

print("Opening image...")
img = np.loadtxt(source, skiprows=6)
print("Image opened")

# 遮盖高程低于70m的区域,低于70米赋值为1，高于70米赋值为0
wet = np.where(img < 70, 1, 0)
print("Image masked")

# 使用内置的linecache模块循环解析头部信息
hdr = [getline(source, i) for i in range(1, 7)]
values = [float(h.split(" ")[-1].strip()) for h in hdr]
cols, rows, lx, ly, cell, nd = values  # 取出图片每个像素的对应的跨度cell
xres = cell
yres = cell * -1

# 屏幕坐标中洪水淹没的起始点
sx = 2582
sy = 2057
print("Beginning flood fill ...")
fld = floodFill(sx, sy, wet)
print("Finished flood fill ...")

header = ""
for i in range(6):
    header += hdr[i]

print("Saving grid")
# 打开输出文件，添加hdr值，保存数组
with open(output, "wb") as f:
    f.write(bytes(header, "UTF-8"))
    np.savetxt(f, fld, fmt="%li")
print("Done!")

# 将asc转换为图片


